import WordMetrics
import numpy as np
from dtwalign import dtw_from_distance_matrix
import time
from typing import List, Tuple
import re

WORD_NOT_FOUND_TOKEN = '-'

def get_word_distance_matrix(words_estimated: List[str], words_real: List[str]) -> np.ndarray:
    words_estimated = words_estimated if words_estimated is not None else []
    words_real = words_real if words_real is not None else []

    number_of_real_words = len(words_real)
    number_of_estimated_words = len(words_estimated)
    
    offset_blank = 1
    matrix_rows = number_of_estimated_words + offset_blank
    matrix_cols = number_of_real_words
    word_distance_matrix = np.full((matrix_rows, matrix_cols), np.inf)

    for idx_estimated in range(number_of_estimated_words):
        for idx_real in range(number_of_real_words):
            cost = WordMetrics.edit_distance_python(
                words_estimated[idx_estimated], words_real[idx_real]
            )
            word_distance_matrix[idx_estimated, idx_real] = cost

    if offset_blank == 1:
        for idx_real in range(number_of_real_words):
            word_distance_matrix[number_of_estimated_words, idx_real] = len(words_real[idx_real]) 
    return word_distance_matrix

def get_resulting_string_from_dtw_path(dtw_alignment_path: List[Tuple[int, int]], 
                                     words_estimated: List[str], 
                                     words_real: List[str]) -> Tuple[List[str], List[int]]:
    if not isinstance(dtw_alignment_path, (list, tuple, np.ndarray)) or \
       (len(dtw_alignment_path) > 0 and not (isinstance(dtw_alignment_path[0], (list, tuple, np.ndarray)) and len(dtw_alignment_path[0]) == 2)):
        print("[ERROR] DTW路径格式异常，全部标记为遗漏")
        mapped_words = [WORD_NOT_FOUND_TOKEN] * len(words_real)
        mapped_words_indices = [-1] * len(words_real)
        return mapped_words, mapped_words_indices

    mapped_words = [WORD_NOT_FOUND_TOKEN] * len(words_real)
    mapped_words_indices = [-1] * len(words_real)

    for real_idx, est_idx in dtw_alignment_path:
        if real_idx < len(words_real):
            if est_idx < len(words_estimated):
                if mapped_words[real_idx] == WORD_NOT_FOUND_TOKEN:
                    mapped_words[real_idx] = words_estimated[est_idx]
                    mapped_words_indices[real_idx] = est_idx
    return mapped_words, mapped_words_indices

def lcs_align(words_estimated, words_real):
    m, n = len(words_estimated), len(words_real)
    dp = [[0] * (n + 1) for _ in range(m + 1)]

    for i in range(1, m + 1):
        for j in range(1, n + 1):
            if words_estimated[i - 1] == words_real[j - 1]:
                dp[i][j] = dp[i - 1][j - 1] + 1
            else:
                dp[i][j] = max(dp[i - 1][j], dp[i][j - 1])
    
    mapped_words = [WORD_NOT_FOUND_TOKEN] * n
    mapped_indices = [-1] * n
    
    i, j = m, n
    while i > 0 and j > 0:
        if words_estimated[i - 1] == words_real[j - 1]:
            mapped_words[j - 1] = words_estimated[i - 1]
            mapped_indices[j-1] = i - 1
            i -= 1
            j -= 1
        elif dp[i - 1][j] > dp[i][j - 1]:
            i -= 1
        else:
            j -= 1
            
    return mapped_words, mapped_indices

def get_best_mapped_words(words_estimated: List[str], words_real: List[str]) -> Tuple[List[str], List[int]]:
    if not words_real:
        return [], [] 
    if not words_estimated:
        return [WORD_NOT_FOUND_TOKEN] * len(words_real), [-1] * len(words_real)

    word_distance_matrix = get_word_distance_matrix(words_estimated, words_real)

    try:
        calculated_window_size = max(1,int(0.2*max(len(words_real),len(words_estimated))))
        
        if len(words_real) == 1:
            if word_distance_matrix.T.shape[1] > 0:
                best_est_idx = np.argmin(word_distance_matrix.T[0])
                dtw_path = [(0, int(best_est_idx))]
            else:
                dtw_path = []
        elif word_distance_matrix.shape[1] == 0 or word_distance_matrix.shape[0] == 0:
            dtw_path = []
        else:
            alignment = dtw_from_distance_matrix(word_distance_matrix.T, window_type="sakoechiba", window_size=calculated_window_size)
            dtw_path = alignment.get_warping_path()
            if not isinstance(dtw_path, (list, tuple, np.ndarray)) or \
               (len(dtw_path) > 0 and not (isinstance(dtw_path[0], (list, tuple, np.ndarray)) and len(dtw_path[0]) == 2)):
                raise ValueError("DTW路径格式异常")

        mapped_words, mapped_words_indices = get_resulting_string_from_dtw_path(
            dtw_path, words_estimated, words_real
        )
    except Exception as e:
        print(f"[ERROR] DTW对齐失败: {e}，尝试LCS降级对齐")
        mapped_words, mapped_words_indices = lcs_align(words_estimated, words_real)

    return mapped_words, mapped_words_indices 